package utils

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"sort"
	"strings"
)

// WecomCrypto 企业微信加解密工具
type WecomCrypto struct {
	Token          string
	EncodingAESKey string
	CorpID         string
	aesKey         []byte
}

// NewWecomCrypto 创建企业微信加解密器
func NewWecomCrypto(token, encodingAESKey, corpID string) *WecomCrypto {
	// encodingAESKey 是 base64 编码的，需要解码
	key, _ := base64.StdEncoding.DecodeString(encodingAESKey + "=")
	return &WecomCrypto{
		Token:          token,
		EncodingAESKey: encodingAESKey,
		CorpID:         corpID,
		aesKey:         key,
	}
}

// ValidateSignature 验证回调签名
func (w *WecomCrypto) ValidateSignature(token, timestamp, nonce, echoStr, msgSignature string) bool {
    signature := w.GenarateSignature(token, timestamp, nonce, echoStr)
    return signature == msgSignature
}
// GenarateSignature 生成签名
func (w *WecomCrypto) GenarateSignature(token, timestamp, nonce, echoStr string) string {
    strs := []string{token, timestamp, nonce, echoStr}
    sort.Strings(strs)
    str := strings.Join(strs, "")
    sum := sha1.Sum([]byte(str))
    return fmt.Sprintf("%x", sum)
}

// DecryptEchoStr 解密 echostr，返回明文
func (w *WecomCrypto) DecryptEchoStr(echoStr string) (string, error) {
	// 1. Base64 解码
	cipherText, err := base64.StdEncoding.DecodeString(echoStr)
	if err != nil {
		return "", err
	}

	// 2. AES-256-CBC 解密
	block, err := aes.NewCipher(w.aesKey)
	if err != nil {
		return "", err
	}

	mode := cipher.NewCBCDecrypter(block, w.aesKey[0:16])
	mode.CryptBlocks(cipherText, cipherText)

	// 3. 去除 PKCS#7 填充
	plaintext := w.pkcs7Unpad(cipherText)
	if len(plaintext) < 16 {
		return "", fmt.Errorf("invalid plaintext")
	}

	// 4. 去除 16 字节的随机数
	// 5. 提取 4 字节的 msg_len (我们这里忽略)
	// 6. 提取 msg (从第 20 字节开始)
	// 7. 校验最后的 CorpID
	if len(plaintext) < 20+len(w.CorpID) {
		return "", fmt.Errorf("invalid plaintext length")
	}

	msgLen := int(plaintext[19]) + int(plaintext[18])<<8 + int(plaintext[17])<<16 + int(plaintext[16])<<24
	msgStart := 20
	msgEnd := msgStart + msgLen
	if msgEnd > len(plaintext) {
		return "", fmt.Errorf("invalid msg length")
	}

	msg := plaintext[msgStart:msgEnd]
	receivedCorpID := plaintext[msgEnd:]
	if string(receivedCorpID) != w.CorpID {
		return "", fmt.Errorf("corp_id mismatch")
	}

	return string(msg), nil
}

// pkcs7Unpad 去除 PKCS#7 填充
func (w *WecomCrypto) pkcs7Unpad(data []byte) []byte {
	padding := int(data[len(data)-1])
	if padding > len(data) {
		return data
	}
	return data[:len(data)-padding]
}